import pandas as pd
import numpy as np
import sys
from pathlib import Path
import chardet

# 设置项目根目录路径
PROJECT_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = PROJECT_ROOT / "data"

def detect_encoding(file_path):
    """检测文件编码"""
    with open(file_path, 'rb') as f:
        raw_data = f.read(10000)
    result = chardet.detect(raw_data)
    return result['encoding'] or 'gbk'

def standardize_column_names(df):
    """标准化列名"""
    column_mapping = {
        '收盘价': ['收盘价', 'close', 'close price', 'Close', '收盘'],
        '开盘价': ['开盘价', 'open', 'open price', 'Open', '开盘'],
        '最高价': ['最高价', 'high', 'high price', 'High', '最高'],
        '最低价': ['最低价', 'low', 'low price', 'Low', '最低'],
        '成交量': ['成交量', 'volume', 'Volume', '成交股数'],
        '成交额': ['成交额', 'turnover', 'Turnover', 'amount', 'Amount', '成交金额'],
        '换手率': ['换手率', 'turnover rate', 'Turnover Rate'],
        '涨跌幅': ['涨跌幅', 'change', 'Change', 'pct_change', '涨跌幅度'],
        '股票代码': ['股票代码', 'stock code', 'Stock Code', 'code', 'symbol'],
        '日期': ['日期', 'date', 'Date', '交易日期']
    }

    new_columns = []
    for col in df.columns:
        col_clean = str(col).strip().replace(" ", "").lower()
        matched = False
        for standard_name, aliases in column_mapping.items():
            for alias in aliases:
                alias_clean = str(alias).strip().replace(" ", "").lower()
                if col_clean == alias_clean:
                    new_columns.append(standard_name)
                    matched = True
                    break
            if matched:
                break
        if not matched:
            new_columns.append(col)

    df.columns = new_columns
    return df


def calculate_features(df):
    """计算关键特征 - 增强版"""
    # 添加基础列检查
    required_base_cols = ['最高价', '最低价', '收盘价', '成交量']
    missing_base = [col for col in required_base_cols if col not in df.columns]
    if missing_base:
        print(f"警告: 缺失基础列{missing_base}，部分特征无法计算")

    # 基本特征
    if '最高价' in df.columns and '最低价' in df.columns and '开盘价' in df.columns:
        df['price_range'] = (df['最高价'] - df['最低价']) / df['开盘价']
    if '收盘价' in df.columns and '开盘价' in df.columns:
        df['price_change'] = df['收盘价'] / df['开盘价'] - 1

    # 移动平均
    if '收盘价' in df.columns:
        df['5d_ma'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(5).mean())
        df['10d_ma'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(10).mean())
        df['20d_ma'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(20).mean())
        # 避免除零错误
        df['ma_ratio_5_10'] = np.where(df['10d_ma'] > 0, df['5d_ma'] / df['10d_ma'], 0)
        df['ma_ratio_5_20'] = np.where(df['20d_ma'] > 0, df['5d_ma'] / df['20d_ma'], 0)

    # MACD指标
    if '收盘价' in df.columns:
        df['12d_ema'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.ewm(span=12).mean())
        df['26d_ema'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.ewm(span=26).mean())
        df['macd'] = df['12d_ema'] - df['26d_ema']
        df['macd_signal'] = df.groupby('股票代码')['macd'].transform(lambda x: x.ewm(span=9).mean())
        df['macd_hist'] = df['macd'] - df['macd_signal']

    # 波动率
    if '收盘价' in df.columns:
        df['5d_std'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(5).std())
        df['10d_std'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(10).std())
        df['20d_std'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.rolling(20).std())

    # 布林带
    if '20d_ma' in df.columns and '20d_std' in df.columns:
        df['bollinger_upper'] = df['20d_ma'] + 2 * df['20d_std']
        df['bollinger_lower'] = df['20d_ma'] - 2 * df['20d_std']

        # 避免除零错误
        bollinger_diff = df['bollinger_upper'] - df['bollinger_lower']
        df['bollinger_pct'] = np.where(
            bollinger_diff > 0,
            (df['收盘价'] - df['bollinger_lower']) / bollinger_diff,
            0.5
        )
        df['bollinger_width'] = np.where(
            df['20d_ma'] > 0,
            bollinger_diff / df['20d_ma'],
            0
        )

    # 成交量特征
    if '成交量' in df.columns:
        df['volume_ma5'] = df.groupby('股票代码')['成交量'].transform(lambda x: x.rolling(5).mean())
        df['volume_ma10'] = df.groupby('股票代码')['成交量'].transform(lambda x: x.rolling(10).mean())
        if 'volume_ma5' in df.columns:
            # 避免除零错误
            df['volume_ratio'] = np.where(
                df['volume_ma5'] > 0,
                df['成交量'] / df['volume_ma5'],
                1
            )

    # 价格动量
    if '收盘价' in df.columns:
        df['momentum_5'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.pct_change(5))
        df['momentum_10'] = df.groupby('股票代码')['收盘价'].transform(lambda x: x.pct_change(10))

    # 相对强弱指标 (RSI)
    if '收盘价' in df.columns:
        delta = df.groupby('股票代码')['收盘价'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.groupby(df['股票代码']).transform(
            lambda x: x.ewm(alpha=1 / 14, min_periods=14).mean())
        avg_loss = loss.groupby(df['股票代码']).transform(
            lambda x: x.ewm(alpha=1 / 14, min_periods=14).mean())
        rs = np.where(avg_loss > 0, avg_gain / avg_loss, np.nan)
        df['rsi'] = 100 - (100 / (1 + rs))

    # 资金流指标 - 修复部分
    if all(col in df.columns for col in ['最高价', '最低价', '收盘价', '成交量']):
        # 先创建money_flow列
        df['money_flow'] = (df['最高价'] + df['最低价'] + df['收盘价']) / 3 * df['成交量']

        # 然后计算移动平均
        df['money_flow_ma5'] = df.groupby('股票代码')['money_flow'].transform(
            lambda x: x.rolling(5).mean()
        )

    # 在返回前添加保护层
    for col in df.columns:
        if df[col].dtype.kind in 'f':  # 只处理浮点数列
            # 替换无穷大
            df[col] = df[col].replace([np.inf, -np.inf], np.nan)

            # 基于分位数限制极值
            q01 = df[col].quantile(0.01)
            q99 = df[col].quantile(0.99)
            df[col] = np.where(
                df[col] < q01, q01,
                np.where(df[col] > q99, q99, df[col])
            )

    return df.fillna(0)


def prepare_data(predict=False):
    """准备训练和测试数据"""
    DATA_DIR.mkdir(parents=True, exist_ok=True)

    # 读取训练和测试数据
    train_path = DATA_DIR / "train.csv"
    test_path = DATA_DIR / "test.csv"

    # 读取训练数据
    train_encoding = detect_encoding(train_path)
    train = pd.read_csv(train_path, encoding=train_encoding)
    train = standardize_column_names(train)
    print(f"训练数据行数: {len(train)}, 列名: {list(train.columns)}")

    # 读取测试数据
    if test_path.exists():
        test_encoding = detect_encoding(test_path)
        test = pd.read_csv(test_path, encoding=test_encoding)
        test = standardize_column_names(test)
        print(f"测试数据行数: {len(test)}, 列名: {list(test.columns)}")
        full_df = pd.concat([train, test])
    else:
        full_df = train.copy()
        print("未找到测试数据，仅使用训练数据")

    # 计算特征
    full_df = calculate_features(full_df)
    # 创建目标变量
    if '涨跌幅' in full_df.columns:
        print("使用'涨跌幅'列创建目标变量")
        full_df['target'] = full_df.groupby('股票代码')['涨跌幅'].shift(-1)
    elif '收盘价' in full_df.columns:
        print("使用'收盘价'列创建目标变量")
        full_df['target'] = full_df.groupby('股票代码')['收盘价'].pct_change().shift(-1)
    else:
        print("错误: 找不到'涨跌幅'或'收盘价'列，无法创建目标变量")
        sys.exit(1)
    # 关键特征列表
    features = [
        '开盘价', '收盘价', '最高价', '最低价', '成交量', '成交额', '换手率','price_range', 'price_change', '5d_ma', '10d_ma', '20d_ma', 'ma_ratio_5_10', 'ma_ratio_5_20', '5d_std', '10d_std', '20d_std', 'volume_ma5', 'volume_ma10', 'volume_ratio', 'momentum_5', 'momentum_10','rsi', 'macd', 'macd_signal', 'macd_hist', 'bollinger_pct', 'bollinger_width','money_flow_ma5'
    ]
    # 只保留存在的特征
    available_features = [feat for feat in features if feat in full_df.columns]
    print(f"可用特征数量: {len(available_features)}")
    # 处理缺失值
    full_df = full_df.dropna(subset=available_features + ['target'])

    # 获取日期信息
    if '日期' in full_df.columns:
        dates = full_df['日期']
    else:
        print("警告: 找不到'日期'列")
        dates = None

    if predict:
        # 获取最新日期的数据
        if dates is not None:
            last_date = dates.max()
            test_df = full_df[dates == last_date]
        else:
            test_df = full_df.tail(300)
        print(f"测试集行数: {len(test_df)}")
        return None, None, test_df[available_features], test_df['股票代码']
    else:
        # 排除最后一天的数据
        if dates is not None:
            train_df = full_df[dates != dates.max()]
        else:
            train_df = full_df.iloc[:-300]  # 排除最后300行作为测试
        print(f"训练集行数: {len(train_df)}")
        return train_df[available_features], train_df['target'], None, None